# coding=utf-8
"""
作者：dnscmd@qq.com
创建时间：2025/3/7 23:13
描述：
密钥文件有两行组成
cat xx.key
salt
key
"""
import base64
import hashlib
import os.path
import time
from random import sample as randomsample

import wx
import wx.xrc
from Crypto.Cipher import AES
from Crypto.Hash import SHA256
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Random import get_random_bytes
from cryptography.fernet import Fernet
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC


class MySec(wx.Frame):

    def __init__(self, parent):
        wx.Frame.__init__(self, parent, id=wx.ID_ANY, title=wx.EmptyString, pos=wx.DefaultPosition,
                          size=wx.Size(550, 564), style=wx.DEFAULT_FRAME_STYLE | wx.TAB_TRAVERSAL)

        self.SetSizeHints(wx.DefaultSize, wx.DefaultSize)

        bSizer1 = wx.BoxSizer(wx.VERTICAL)

        fgSizer1 = wx.FlexGridSizer(0, 2, 0, 0)
        fgSizer1.SetFlexibleDirection(wx.BOTH)
        fgSizer1.SetNonFlexibleGrowMode(wx.FLEX_GROWMODE_SPECIFIED)

        self.btn_load_key_file = wx.Button(self, wx.ID_ANY, u"载入密钥", wx.DefaultPosition, wx.DefaultSize, 0)
        fgSizer1.Add(self.btn_load_key_file, 0, wx.ALL, 5)

        self.txt_key_file = wx.StaticText(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize, 0)
        self.txt_key_file.Wrap(-1)
        fgSizer1.Add(self.txt_key_file, 0, wx.ALL, 5)

        self.btn_dec_key_key = wx.Button(self, wx.ID_ANY, u"解密Key", wx.DefaultPosition, wx.DefaultSize, 0)
        fgSizer1.Add(self.btn_dec_key_key, 0, wx.ALL, 5)
        self.txt_key_pwd = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.Size(410, 30), style=wx.TE_PASSWORD)
        fgSizer1.Add(self.txt_key_pwd, 0, wx.ALL, 5)

        self.m_staticText2 = wx.StaticText(self, wx.ID_ANY, u"源 路径", wx.DefaultPosition, wx.DefaultSize, 0)
        self.m_staticText2.Wrap(-1)
        fgSizer1.Add(self.m_staticText2, 0, wx.ALL, 5)

        self.txt_source_path = wx.StaticText(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize, 0)
        self.txt_source_path.Wrap(-1)
        fgSizer1.Add(self.txt_source_path, 0, wx.ALL, 5)

        self.m_staticText5 = wx.StaticText(self, wx.ID_ANY, u"目标路径", wx.DefaultPosition, wx.DefaultSize, 0)
        self.m_staticText5.Wrap(-1)
        fgSizer1.Add(self.m_staticText5, 0, wx.ALL, 5)

        self.txt_target_path = wx.StaticText(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize, 0)
        self.txt_target_path.Wrap(-1)
        fgSizer1.Add(self.txt_target_path, 0, wx.ALL, 5)

        bSizer1.Add(fgSizer1, 1, wx.EXPAND, 5)

        bSizer2 = wx.BoxSizer(wx.VERTICAL)

        bSizer1.Add(bSizer2, 1, wx.EXPAND, 5)

        gbSizer2 = wx.GridBagSizer(0, 0)
        gbSizer2.SetFlexibleDirection(wx.BOTH)
        gbSizer2.SetNonFlexibleGrowMode(wx.FLEX_GROWMODE_SPECIFIED)

        self.btn_select_target_dir = wx.Button(self, wx.ID_ANY, u"输出位置", wx.DefaultPosition, wx.DefaultSize, 0)
        gbSizer2.Add(self.btn_select_target_dir, wx.GBPosition(0, 2), wx.GBSpan(1, 1), wx.ALL, 5)
        self.btn_select_target_dir.SetToolTip("不选择将覆盖源文件")

        self.btn_select_source_dir = wx.Button(self, wx.ID_ANY, u"源 目录", wx.DefaultPosition, wx.DefaultSize, 0)
        gbSizer2.Add(self.btn_select_source_dir, wx.GBPosition(0, 1), wx.GBSpan(1, 1), wx.ALL, 5)

        self.btn_select_source_file = wx.Button(self, wx.ID_ANY, u"源 文件", wx.DefaultPosition, wx.DefaultSize, 0)
        gbSizer2.Add(self.btn_select_source_file, wx.GBPosition(0, 0), wx.GBSpan(1, 1), wx.ALL, 5)

        bSizer1.Add(gbSizer2, 1, wx.EXPAND, 5)

        gbSizer4 = wx.GridBagSizer(0, 0)
        gbSizer4.SetFlexibleDirection(wx.BOTH)
        gbSizer4.SetNonFlexibleGrowMode(wx.FLEX_GROWMODE_SPECIFIED)

        self.btn_enc = wx.Button(self, wx.ID_ANY, u"加密", wx.DefaultPosition, wx.DefaultSize, 0)
        gbSizer4.Add(self.btn_enc, wx.GBPosition(0, 0), wx.GBSpan(1, 1), wx.ALL, 5)

        self.btn_dec = wx.Button(self, wx.ID_ANY, u"解密", wx.DefaultPosition, wx.DefaultSize, 0)
        gbSizer4.Add(self.btn_dec, wx.GBPosition(0, 1), wx.GBSpan(1, 1), wx.ALL, 5)

        self.btn_gen_key_file = wx.Button(self, wx.ID_ANY, u"生成密钥", wx.DefaultPosition, wx.DefaultSize, 0)
        gbSizer4.Add(self.btn_gen_key_file, wx.GBPosition(0, 2), wx.GBSpan(1, 1), wx.ALL, 5)

        bSizer1.Add(gbSizer4, 1, wx.EXPAND, 5)

        self.txt_out = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.Size(550, 300), wx.TE_MULTILINE)
        bSizer1.Add(self.txt_out, 0, wx.ALL, 5)

        self.SetSizer(bSizer1)
        self.Layout()

        self.Centre(wx.BOTH)

        self.btn_load_key_file.Bind(wx.EVT_BUTTON, self.load_key_file)
        self.btn_select_source_file.Bind(wx.EVT_BUTTON, self.select_source_file)
        self.btn_select_source_dir.Bind(wx.EVT_BUTTON, self.select_source_dir)
        self.btn_select_target_dir.Bind(wx.EVT_BUTTON, self.select_target_dir)
        self.btn_enc.Bind(wx.EVT_BUTTON, self.enc_files)
        self.btn_dec.Bind(wx.EVT_BUTTON, self.dec_files)
        self.btn_gen_key_file.Bind(wx.EVT_BUTTON, self.gen_key_file)
        self.btn_dec_key_key.Bind(wx.EVT_BUTTON, self.dec_key_key)

        # 密码
        self._header = b"dnpenc"  # 文件头标识
        self._salt = b""  # 盐值 从配置文件载入
        self._iterations = 1000000  # PBKDF2迭代次数
        self._key_len = 32  # AES 256位密钥长度
        self._hmac_len = 32  # SHA256的输出长度
        self._pwd = ""  # 密码
        self._forbidden_dirs = {r'c:\windows', r'c:\programdata', r'c:\program files'}

        key_path = r'C:\workspace\01 新强\dnpenc.key'
        if os.path.exists(key_path):
            self.txt_key_file.SetLabelText(key_path)
            self._load_key_file()

    def __del__(self):
        pass

    def generate_key(self):
        return PBKDF2(self._pwd, self._salt, dkLen=self._key_len, count=self._iterations, hmac_hash_module=SHA256)

    def encrypt_file(self, input_file, output_file):
        if input_file.endswith('.key'):  # key文件
            return

        key = self.generate_key()

        with open(input_file, "rb") as f:
            plaintext = f.read()

        # 如果已经加密过了就不再加密了
        if plaintext.startswith(self._header):
            with open(output_file, "wb") as f:
                f.write(plaintext)
            return

        nonce = get_random_bytes(12)  # 生成随机 Nonce
        cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)  # AES-GCM 加密器
        ciphertext, tag = cipher.encrypt_and_digest(plaintext)  # 加密 + 生成认证标签

        with open(output_file, "wb") as f:
            f.write(self._header)  # 写入文件头
            f.write(nonce)  # 写入 AES-GCM Nonce
            f.write(tag)  # 写入 GCM 认证标签
            f.write(ciphertext)  # 写入加密数据

    def decrypt_file(self, input_file, output_file):
        key = self.generate_key()

        with open(input_file, "rb") as f:
            header = f.read(len(self._header))
            if header != self._header:
                return f"文件头不匹配，解密失败：{input_file}"

            nonce = f.read(12)  # 读取 AES-GCM Nonce
            tag = f.read(16)  # 读取 GCM 认证标签
            ciphertext = f.read()  # 读取加密数据

        cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)  # 创建 AES-GCM 解密器
        try:
            plaintext = cipher.decrypt_and_verify(ciphertext, tag)  # 解密并验证

            with open(output_file, "wb") as f:
                f.write(plaintext)  # 写入解密后的文件
        except ValueError:
            return f"❌ 解密失败，密码错误或数据被篡改：{input_file}"
        return ''

    def load_key_file(self, event):
        dialog = wx.FileDialog(self, "选择密钥文件", wildcard="All files (*.key)|*.key", style=wx.FD_OPEN | wx.FD_FILE_MUST_EXIST)

        if dialog.ShowModal() == wx.ID_OK:
            self.txt_key_file.SetLabelText(dialog.GetPath())
        dialog.Destroy()
        self._load_key_file()

    @staticmethod
    def generate_key_key_from_password(password: str):
        salt = hashlib.md5(password.encode('utf-8')).digest()[5:5 + 8]
        kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100000, backend=default_backend())
        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        return key

    def _load_key_file(self):
        try:
            key_path = self.txt_key_file.GetLabelText()
            if not key_path or not os.path.exists(key_path):
                self.txt_out.SetValue("请指定key的路径")
                return

            with open(key_path, 'r') as f:
                lines = f.readlines()

            if not lines:  # 没有数据
                self.txt_out.SetValue("非法密钥文件")
                return

            if len(lines) == 2:  # 有两行数据，认为未加密
                lines = f.readlines()
                self._salt = lines[0][:32].encode('utf8')
                self._pwd = lines[1]
                self.txt_out.SetValue("载入密钥成功")
                return

            # 遗憾数据认为是加密过的密钥，解密密钥
            if len(lines) == 1:
                key_pwd = self.txt_key_pwd.GetValue()
                if not key_pwd or len(key_pwd) < 8:
                    self.txt_out.SetValue("密钥太短")
                    return

                # 构造解密密钥
                lines = Fernet(MySec.generate_key_key_from_password(key_pwd)).decrypt(lines[0].encode('utf8')).decode().split('\n')
                self._salt = lines[0][:32].encode('utf8')
                self._pwd = lines[1]
                self.txt_out.SetValue("载入密钥成功")
                return
            self.txt_out.SetValue("非法密钥文件")
        except Exception as e:
            self.txt_out.SetValue(e)

    def select_source_file(self, event):
        dialog = wx.FileDialog(self, "选择源文件", wildcard="All files (*.*)|*.*", style=wx.FD_OPEN | wx.FD_FILE_MUST_EXIST)

        if dialog.ShowModal() == wx.ID_OK:
            self.txt_source_path.SetLabelText(dialog.GetPath())
        dialog.Destroy()

    def select_source_dir(self, event):
        dialog = wx.DirDialog(self, "选择源目录", style=wx.DD_DEFAULT_STYLE)

        if dialog.ShowModal() == wx.ID_OK:
            self.txt_source_path.SetLabelText(dialog.GetPath())
        dialog.Destroy()

    def select_target_dir(self, event):
        dialog = wx.DirDialog(self, "选择保存位置", style=wx.DD_DEFAULT_STYLE)

        if dialog.ShowModal() == wx.ID_OK:
            self.txt_target_path.SetLabelText(dialog.GetPath())
        dialog.Destroy()

    def gen_key_file(self, event):
        def get_random_chars(length=32):
            ls = randomsample(r'qwertyuiopasdfghjklzxcvbnmQWERTYUIOPASDFGHJKLZXCVBNM0987654321~!@#$%^&*()_+<>?:{},./;\|', length)
            return ''.join(ls)

        key, salt = get_random_chars(), get_random_chars()
        with open(f'{int(time.time())}.key', 'w') as f:
            f.writelines([salt, '\n', key])

    @staticmethod
    def enc_key(pwd: str, txt: str):
        if not pwd or len(pwd) < 8:
            print('密码不应小于8位')
            return
        print(Fernet(MySec.generate_key_key_from_password(pwd)).encrypt(txt.encode()).decode())

    def dec_key_key(self, event):
        key_pwd = self.txt_key_pwd.GetValue()
        if not key_pwd or len(key_pwd) < 8:
            self.txt_out.SetValue("密钥太短")
            return
        self._load_key_file()

    def _get_files(self):

        path = self.txt_source_path.GetLabelText()
        file_ls = []

        if os.path.isfile(path):
            file_ls.append(path)
        elif os.path.isdir(path):
            for root, _, fnames in os.walk(path):
                for f in fnames:
                    file_ls.append(os.path.join(root, f))

        return file_ls

    @staticmethod
    def gen_out_path(target_dir: str, source_dir: str, input_full_path: str, is_enc=True) -> str:
        """
        生成输出文件路径和文件夹路径
        :param is_enc: 是否是加密
        :param source_dir: 用户选择的源文件夹
        :param target_dir: 目标文件夹
        :param input_full_path: 输入文件夹
        :return: 输出路径
        """
        if not target_dir:
            return input_full_path

        source_dir = source_dir.lower()
        input_full_path = os.path.join(os.path.dirname(input_full_path).lower(), os.path.basename(input_full_path))

        # 用户选择的就是个文件
        if source_dir == input_full_path:
            return os.path.join(target_dir, os.path.basename(input_full_path))

        # 用户选的是一个文件夹
        # 如果源目录是 c:\usr\xx   目标目录为d:\download ,拿实际的目标文件夹应该为 d:\download\xx
        if is_enc:
            out_path = input_full_path.replace(source_dir, os.path.join(target_dir, f'{os.path.basename(source_dir)}_已加密'))
        else:
            out_path = input_full_path.replace(source_dir, os.path.join(target_dir, f'{os.path.basename(source_dir)}'))

        d = os.path.dirname(out_path)
        if not os.path.exists(d):
            os.makedirs(d)
        return out_path

    def enc_files(self, event):
        self.txt_out.SetValue('')

        source_path = self.txt_source_path.GetLabelText().lower()
        # 排除禁用路径
        if source_path.startswith('c:') and len(source_path) < 10:
            self.txt_out.SetValue("安全考虑，不允许加密该路径")
            return

        for fb in self._forbidden_dirs:
            if source_path.startswith(fb):
                self.txt_out.SetValue("安全考虑，不允许加密该路径")
                return

        target_dir = self.txt_target_path.GetLabelText()

        dialog = None
        fail_ls = []

        try:
            file_ls = self._get_files()
            if not file_ls:
                return

            file_num = len(file_ls)

            dialog = wx.ProgressDialog("进度条", "正在执行任务...", maximum=len(file_ls), style=wx.PD_ELAPSED_TIME)
            for i in range(file_num):
                try:
                    out_file = MySec.gen_out_path(target_dir, source_path, file_ls[i])
                    self.encrypt_file(file_ls[i], out_file)
                except Exception as ei:
                    fail_ls.append(str(ei))
                dialog.Update(i + 1, f'完成 {(i + 1)}/{file_num}')
        except Exception as e:
            self.txt_out.SetValue(str(e))

        if fail_ls:
            self.txt_out.SetValue('\n'.join(fail_ls))
        else:
            self.txt_out.SetValue("加密完成")

        if dialog is not None:
            dialog.Close()
            dialog.Destroy()

    def dec_files(self, event):
        self.txt_out.SetValue('')
        target_dir = self.txt_target_path.GetLabelText()

        dialog = None
        try:
            file_ls = self._get_files()
            if not file_ls:
                return
            source_path = self.txt_source_path.GetLabelText().lower()
            file_num = len(file_ls)
            dialog = wx.ProgressDialog("进度条", "正在执行任务...", maximum=len(file_ls), style=wx.PD_ELAPSED_TIME)
            fail_ls = []
            for i in range(file_num):
                try:
                    out_file = MySec.gen_out_path(target_dir, source_path, file_ls[i], is_enc=False)
                    res = self.decrypt_file(file_ls[i], out_file)
                    if res:
                        fail_ls.append(res)
                except Exception as ei:
                    fail_ls.append(str(ei))
                dialog.Update(i + 1, f'完成 {(i + 1)}/{file_num}')
            if fail_ls:
                self.txt_out.SetValue('\n'.join(fail_ls))
            else:
                self.txt_out.SetValue("解密完成")
        except Exception as e:
            self.txt_out.SetValue(str(e))

        if dialog is not None:
            dialog.Close()
            dialog.Destroy()


if __name__ == '__main__':
    app = wx.App()
    frame = MySec(None)
    frame.Show()
    app.MainLoop()

    # with open(r'A:\01 新强\01 dnp.key','r') as f:
    #     c = f.read()
    #     MySec.enc_key('DuuDhy5211314$$', c)
